"""
WebSocket聊天功能模块
提供用户和客服之间的实时聊天功能
"""

from typing import Dict, List
from fastapi import WebSocket
from starlette.websockets import WebSocketDisconnect
import json
from datetime import datetime


# ============ 全局连接管理 ============
# 
# 全局连接管理字典，用于维护所有活跃的WebSocket连接
# 这些连接在内存中存储，应用重启后会丢失（生产环境建议使用Redis等持久化方案）

# 客服连接管理字典
# 键: 客服ID（service_id），值: 客服的WebSocket连接对象
# 用途: 存储所有在线客服的连接，用于接收咨询请求通知
service_connections: Dict[str, WebSocket] = {}

# 用户连接管理字典
# 键: 用户ID（user_id），值: 用户的WebSocket连接对象
# 用途: 存储所有在线用户的连接，用于接收咨询状态更新通知
user_connections: Dict[str, WebSocket] = {}

# 聊天室管理字典
# 键: 聊天室ID（room_id），格式为 "{user_id}_{service_id}"
# 值: 包含用户和客服WebSocket连接的字典，格式为 {"user": websocket, "service": websocket}
# 用途: 存储每个聊天室的双向连接，实现用户和客服的实时聊天
# 示例: {"user_123_service_1001": {"user": <WebSocket>, "service": <WebSocket>}}
chat_rooms: Dict[str, Dict[str, WebSocket]] = {}

# 咨询请求列表字典
# 键: 客服ID（service_id）
# 值: 该客服的所有咨询请求列表，每个请求包含:
#     - user_id: 用户ID
#     - service_id: 客服ID
#     - created_at: 创建时间
#     - room_id: 聊天室ID
#     - status: 状态（pending/accepted/rejected）
# 用途: 管理咨询请求队列，前端可以通过HTTP接口或WebSocket获取待处理请求
consultation_requests: Dict[str, List[Dict]] = {}


# ============ 工具函数 ============

async def send_json(websocket: WebSocket, data: dict):
    """
    发送JSON格式的消息到WebSocket连接
    
    参数:
        websocket: 目标WebSocket连接对象
        data: 要发送的数据字典，会自动转换为JSON字符串
    
    说明:
        - 使用 ensure_ascii=False 保持中文字符正常显示
        - 异步函数，需要使用 await 调用
    """
    await websocket.send_text(json.dumps(data, ensure_ascii=False))  # 将字典转换为JSON字符串并发送，ensure_ascii=False保持中文正常显示


def get_room_id(user_id: str, service_id: str) -> str:
    """
    生成聊天室唯一ID
    参数:
        user_id: 用户ID
        service_id: 客服ID
    
    返回:
        聊天室ID，格式: "{user_id}_{service_id}"
    
    说明:
        - 同一用户和客服的聊天室ID是固定的
        - 房间ID确保每个用户-客服对都有唯一的聊天空间
    """
    return f"{user_id}_{service_id}"  # 返回格式化的聊天室ID，格式: user_id_service_id


def match_service(user_id: str) -> str:
    """
    为用户匹配可用的客服
    
    参数:
        user_id: 用户ID
    
    返回:
        匹配到的客服ID
    
    说明:
        - 优先返回在线客服
        - 如果没有在线客服，返回默认客服ID
    """
    # 如果有在线客服，返回第一个
    if service_connections:
        return list(service_connections.keys())[0]  # 返回第一个在线客服的ID
    # 如果没有在线客服，返回默认客服ID
    return "service_1001"  # 返回默认客服ID（当没有在线客服时）


# ============ WebSocket端点 ============

async def websocket_chat_consumer(websocket: WebSocket, group: str):
    """
    WebSocket聊天连接端点
    
    连接路径: ws://host/ws/chat/{group}
    其中 group 格式: "user{user_id}" 或 "service{service_id}"
    例如: /ws/chat/user123 或 /ws/chat/service1001
    
    功能说明:
    1. 用户连接时自动创建咨询请求并推送
    2. 客服连接时自动推送所有待处理请求
    3. 实时消息转发（用户↔客服）
    4. 状态更新实时推送
    """
    await websocket.accept()
    
    print(f"[WebSocket连接] 收到连接请求，group: {group}")
    
    # 判断是用户还是客服
    def process_string(input_string, prefix):
        """从字符串中提取指定前缀后的内容"""
        index = input_string.find(prefix)
        if index != -1:
            return input_string[index + len(prefix):]
        return None
    
    userid = process_string(group, 'user')
    serviceid = process_string(group, 'service')
    
    print(f"[WebSocket连接] 解析结果 - userid: {userid}, serviceid: {serviceid}")
    
    try:
        if userid:
            # ============ 用户连接处理 ============
            user_id = f"user_{userid}"
            
            # 存储用户连接
            if user_id not in user_connections:
                user_connections[user_id] = websocket
            print(f"[用户连接] {user_id} 已连接，当前在线用户数：{len(user_connections)}")
            
            # 自动匹配客服
            service_id = match_service(user_id)
            room_id = get_room_id(user_id, service_id)
            
            # 创建咨询请求信息
            request_info = {
                "user_id": user_id,
                "service_id": service_id,
                "room_id": room_id,
                "created_at": datetime.now().isoformat(),
                "status": "pending"
            }
            
            # 添加到咨询请求列表
            if service_id not in consultation_requests:
                consultation_requests[service_id] = []
            
            # 检查是否已存在相同的请求
            existing = [r for r in consultation_requests[service_id] 
                       if r["user_id"] == user_id and r["status"] == "pending"]
            
            if not existing:
                consultation_requests[service_id].append(request_info)
                
                # 实时推送：通知客服有新咨询请求
                if service_id in service_connections:
                    await send_json(service_connections[service_id], {
                        "type": "new_consultation",
                        "data": request_info
                    })
                    print(f"[实时推送] 已向客服 {service_id} 推送新咨询请求：用户 {user_id}")
            
            # 实时推送：通知用户咨询请求已创建
            await send_json(websocket, {
                "type": "consultation_created",
                "code": 200,
                "message": "咨询请求已创建",
                "data": existing[0] if existing else request_info
            })
            
            # 进入消息循环
            while True:
                try:
                    data = await websocket.receive_text()
                    message = json.loads(data)
                    message_type = message.get('type', 'chat')
                    
                    print(f"[用户消息] {user_id} 收到消息类型: {message_type}, 内容: {message}")
                    
                    if message_type == 'chat':
                        # 聊天消息处理
                        content = message.get('content', '')
                        sender = message.get('sender', user_id)
                        timestamp = datetime.now().isoformat()
                        
                        chat_message = {
                            "type": "chat",
                            "content": content,
                            "sender": sender,
                            "timestamp": timestamp,
                            "room_id": room_id,
                        }
                        
                        # 实时推送：转发消息给对方
                        if service_id in service_connections:
                            await send_json(service_connections[service_id], chat_message)
                            print(f"[实时推送] 用户 {sender} 的消息已实时推送给客服，内容: {content}")
                        else:
                            print(f"[警告] 客服 {service_id} 未连接，无法推送消息")
                    
                    elif message_type == 'start_consult':
                        # 手动创建咨询请求（备用）
                        # 逻辑已在连接时处理，这里可以忽略或重复处理
                        pass
                    else:
                        print(f"[用户消息] 未知消息类型: {message_type}")
                        
                except WebSocketDisconnect:
                    # 客户端断开连接，退出循环
                    print(f"[用户断开] {user_id} 连接已断开，退出消息循环")
                    # 清理连接
                    if user_id in user_connections:
                        user_connections.pop(user_id, None)
                    break
                except json.JSONDecodeError as e:
                    print(f"[错误] JSON解析失败: {e}, 原始数据: {data}")
                except Exception as e:
                    print(f"[错误] 处理用户消息失败: {e}")
                    import traceback
                    traceback.print_exc()
                    
        elif serviceid:
            # ============ 客服连接处理 ============
            service_id = f"service_{serviceid}"
            
            print(f"[客服连接] 开始处理客服连接，service_id: {service_id}")
            
            # 存储客服连接
            if service_id not in service_connections:
                service_connections[service_id] = websocket
            else:
                print(f"[警告] 客服 {service_id} 已有连接，覆盖旧连接")
                service_connections[service_id] = websocket
            consultation_requests[service_id] = consultation_requests.get(service_id, [])
            
            print(f"[客服连接] {service_id} 已连接，当前在线客服数：{len(service_connections)}")
            
            # 实时推送：自动推送所有待处理请求
            pending_requests = [r for r in consultation_requests.get(service_id, []) 
                              if r["status"] == "pending"]
            if pending_requests:
                await send_json(websocket, {
                    "type": "consultation_list",
                    "data": pending_requests
                })
                print(f"[实时推送] 已向客服 {service_id} 推送 {len(pending_requests)} 个待处理咨询请求")
            else:
                print(f"[客服连接] 客服 {service_id} 暂无待处理请求")
            
            print(f"[客服连接] 客服 {service_id} 进入消息循环，等待消息...")
            # 进入消息循环
            while True:
                try:
                    data = await websocket.receive_text()
                    message = json.loads(data)
                    message_type = message.get('type', 'chat')
                    print(f"[客服消息] {service_id} 收到消息类型: {message_type}, 内容: {message}")
                    if message_type == 'chat':
                        # 聊天消息处理
                        content = message.get('content', '')
                        sender = message.get('sender', service_id)
                        timestamp = datetime.now().isoformat()
                        # 从room_id或消息中提取用户ID和房间ID
                        room_id = message.get('room_id', '')
                        # 如果没有room_id，尝试从当前客服的咨询请求中查找
                        if not room_id:
                            # 查找当前客服的第一个待处理或已接受的咨询请求
                            for req in consultation_requests.get(service_id, []):
                                if req["status"] in ["pending", "accepted"]:
                                    room_id = req["room_id"]
                                    break
                        if room_id:
                            chat_message = {
                                "type": "chat",
                                "content": content,
                                "sender": sender,
                                "timestamp": timestamp,
                                "room_id": room_id,
                            }
                            # 实时推送：转发消息给用户
                            # 从room_id中提取user_id（格式：user_123_service_1001）
                            parts = room_id.split('_')
                            if len(parts) >= 2:
                                user_id = '_'.join(parts[:2]) if parts[0] == 'user' else parts[0]  # 提取 user_123
                                if user_id in user_connections:
                                    await send_json(user_connections[user_id], chat_message)
                                    print(f"[实时推送] 客服 {sender} 的消息已实时推送给用户 {user_id}，内容: {content}")
                                else:
                                    print(f"[警告] 用户 {user_id} 未连接，无法推送消息")
                            else:
                                print(f"[警告] 无法从room_id {room_id} 中提取用户ID")
                        else:
                            print(f"[警告] 无法确定房间ID，消息无法转发")
                    elif message_type == 'accept_consultation':
                        # 接受咨询请求
                        room_id = message.get('room_id')
                        # 更新请求状态
                        for req in consultation_requests.get(service_id, []):
                            if req["room_id"] == room_id:
                                req["status"] = "accepted"
                                # 实时推送：通知用户咨询已被接受
                                user_id = req["user_id"]
                                if user_id in user_connections:
                                    await send_json(user_connections[user_id], {
                                        "type": "consultation_accepted",
                                        "data": {"room_id": room_id, "service_id": service_id}
                                    })
                                    print(f"[实时推送] 已向用户 {user_id} 推送咨询已被客服 {service_id} 接受")
                                break
                    else:
                        print(f"[客服消息] 未知消息类型: {message_type}")
                        
                except WebSocketDisconnect:
                    # 客户端断开连接，退出循环
                    print(f"[客服断开] {service_id} 连接已断开，退出消息循环")
                    # 清理连接
                    if service_id in service_connections:
                        service_connections.pop(service_id, None)
                    break
                except json.JSONDecodeError as e:
                    print(f"[错误] JSON解析失败: {e}, 原始数据: {data}")
                except Exception as e:
                    print(f"[错误] 处理客服消息失败: {e}")
                    import traceback
                    traceback.print_exc()
        else:
            # 无效的连接参数
            print(f"[错误] 无效的连接参数: {group}, userid={userid}, serviceid={serviceid}")
            print(f"[错误] 连接将被关闭")
            await websocket.close(code=1008, reason="Invalid group format")
            return
            
    except WebSocketDisconnect:
        # 正常断开连接
        if userid:
            user_id = f"user_{userid}"
            if user_id in user_connections:
                user_connections.pop(user_id, None)
                print(f"[用户断开] {user_id} 已断开，当前在线用户数：{len(user_connections)}")
        elif serviceid:
            service_id = f"service_{serviceid}"
            if service_id in service_connections:
                service_connections.pop(service_id, None)
                print(f"[客服断开] {service_id} 已断开，当前在线客服数：{len(service_connections)}")
    except Exception as e:
        print(f"[错误] WebSocket处理异常: {e}")
        import traceback
        traceback.print_exc()
        # 异常时也清理连接
        if userid:
            user_id = f"user_{userid}"
            if user_id in user_connections:
                user_connections.pop(user_id, None)
        elif serviceid:
            service_id = f"service_{serviceid}"
            if service_id in service_connections:
                service_connections.pop(service_id, None)

